## GANET V2
## Copyright (c) 2014 Ganet Development Team. All rights reserved.
##
## This file is subject to the terms and conditions defined in file
## 'LICENSE.md', which is part of this source code package.
## See 'AUTHORS' file for a list of contributors.

import sys, socket, select, time

class Clock(object):

    def __init__(self):
        self.prev_time = time.time()
        self.dt = 0.0

    def tick(self):
        now = time.time()
        self.dt = now - self.prev_time
        self.prev_time = now
        return self.dt

class UdpProxy(object):

    TIMEOUT = 1.0

    def __init__(self, remote_addr, port = None, buffer_size = 4096):
        self._remote = remote_addr
        if not port:
            self._port = self._remote[1]
        else:
            self._port = port
        self._socks = []
        self._main_sock = None
        self._buffer_size = buffer_size
        self._map = {}
        self._invmap = {}

    @property
    def remove_addr(self):
        return self._remote

    @property
    def port(self):
        return self._port

    def run(self):
        self._main_sock = socket.socket(type = socket.SOCK_DGRAM)
        self._main_sock.bind(('', self._port))
        self._socks.append(self._main_sock)

        try:
            cl = Clock()
            while True:
                self.update(cl.tick())
                self._process()
        except KeyboardInterrupt:
            self._cleanup()

    def _log_debug(self, msg):
        #print '[DEBUG] %s' % msg
        pass

    def _log_info(self, msg):
        print '[INFO] %s' % msg

    def send(self, socket, addr, msg):
        socket.sendto(msg, addr)

    def update(self, dt):
        # override this methode to realize special features like
        # delaying messages
        pass

    def _process_local(self):
        try:
            data, addr = self._main_sock.recvfrom(self._buffer_size)
            sock = self._map.get(addr, None)

            if not sock:
                sock = socket.socket(type = socket.SOCK_DGRAM)
                sock.bind(('', 0))
                self._map[addr] = sock
                self._invmap[sock] = addr
                self._socks.append(sock)
                self._log_info('forwarding new client %s:%d' % addr)

            self._log_debug('forwarding message from %s:%d' % addr)
            self.send(sock, self._remote, data)
        except socket.error as e:
            if e.errno == 10054:
                # windows specific
                self._log_info('a client has probably shut down')
            else:
                log_txt = 'got socket error while forwarding to %s%d:'
                log_txt = log_txt % self._remote
                self._log_info(log_txt)

    def _process_remote(self, sock):
        try:
            data = sock.recv(self._buffer_size)
            addr = self._invmap[sock]
            self._log_debug('forwarding message from %s:%d' % addr)
            self.send(self._main_sock, addr, data)
        except socket.error as e:
            if e.errno == 10054:
                # windows specific
                log_txt = 'remote (%s:%d) has probably shut down'
                self._log_info(log_txt % self._remote)
            else:
                log_txt = 'got socket while forwarding to %s%d:'
                self._log_info(log_txt % addr)
                
    def _process(self):
        r, w, e = select.select(self._socks, [], [], 0)

        for sock in r:
            if sock == self._main_sock:
                self._process_local()
            else:
                self._process_remote(sock)

    def _cleanup(self):
        for sock in self._socks:
            sock.close()
        self._socks[:] = []
        self._map.clear()
        self._invmap.clear()

if __name__ == '__main__':

    remote_addr = ('quantumreboot.com', 54321)
    port = 12345
    proxy = UdpProxy(remote_addr, port)
    print 'starting forwarding to %s:%d' % remote_addr
    print 'listening on port %d' % port
    proxy.run()
    print 'forwarding to %s:%d stopped' % remote_addr
